"""
强化学习排班算法实现模块

通过学习历史排班数据，优化未来排班决策。
"""

import random
import logging
from datetime import datetime
from collections import defaultdict


class RLScheduler:
    """强化学习排班算法实现"""

    def __init__(
        self,
        learning_rate=0.1,
        discount_factor=0.9,
        exploration_rate=0.1,
        exploration_decay=0.995,
        min_exploration_rate=0.01,
    ):
        """初始化强化学习排班器

        Args:
            learning_rate: 学习率
            discount_factor: 折扣因子
            exploration_rate: 探索率
            exploration_decay: 探索率衰减
            min_exploration_rate: 最小探索率
        """
        self.logger = logging.getLogger(__name__)
        self.learning_rate = learning_rate
        self.discount_factor = discount_factor
        self.exploration_rate = exploration_rate
        self.exploration_decay = exploration_decay
        self.min_exploration_rate = min_exploration_rate

        # Q值表，用于存储状态-动作对的价值
        self.q_table = defaultdict(lambda: defaultdict(float))

        # 历史数据，用于训练
        self.history = []

    def train(
        self, historical_schedules, caregivers, elders, shift_config, episodes=1000
    ):
        """使用历史排班数据训练Q值表

        Args:
            historical_schedules: 历史排班数据
            caregivers: 护理员列表
            elders: 老人列表
            shift_config: 班次配置
            episodes: 训练轮数

        Returns:
            训练后的Q值表
        """
        self.logger.info(
            f"开始使用历史排班数据训练Q值表，共 {len(historical_schedules)} 条历史数据，训练 {episodes} 轮..."
        )
        start_time = datetime.now()

        # 将历史数据转换为状态-动作-奖励-下一状态序列
        self.history = self.preprocess_historical_data(
            historical_schedules, caregivers, elders, shift_config
        )

        # Q-learning训练
        for episode in range(episodes):
            # 随机选择一个历史数据作为起点
            idx = random.randint(0, len(self.history) - 1)
            state, action, reward, next_state = self.history[idx]

            # 更新Q值
            self.update_q_value(state, action, reward, next_state)

            # 衰减探索率
            self.exploration_rate = max(
                self.min_exploration_rate,
                self.exploration_rate * self.exploration_decay,
            )

            # 每100轮输出一次日志
            if episode % 100 == 0:
                self.logger.info(
                    f"第 {episode} 轮训练完成，当前探索率: {self.exploration_rate:.4f}"
                )

        end_time = datetime.now()
        training_time = (end_time - start_time).total_seconds()
        self.logger.info(
            f"Q值表训练完成，耗时 {training_time:.2f} 秒，Q值表大小: {len(self.q_table)}"
        )

        return self.q_table

    def generate_schedule(
        self, initial_state, caregivers, elders, shift_config, max_steps=100
    ):
        """使用训练好的Q值表生成排班方案

        Args:
            initial_state: 初始状态
            caregivers: 护理员列表
            elders: 老人列表
            shift_config: 班次配置
            max_steps: 最大步数

        Returns:
            生成的排班方案
        """
        self.logger.info("开始使用强化学习算法生成排班方案...")
        start_time = datetime.now()

        state = initial_state
        schedule = self.state_to_schedule(state, caregivers, elders, shift_config)

        for step in range(max_steps):
            # 选择动作
            action = self.select_action(state)

            # 执行动作，获取下一个状态和奖励
            next_state, reward = self.take_action(
                state, action, caregivers, elders, shift_config
            )

            # 更新Q值
            self.update_q_value(state, action, reward, next_state)

            # 更新状态
            state = next_state

            # 更新排班方案
            schedule = self.state_to_schedule(state, caregivers, elders, shift_config)

            # 如果达到终止条件，则结束
            if self.is_terminal_state(state, caregivers, elders, shift_config):
                self.logger.info(f"第 {step} 步达到终止条件，提前结束生成")
                break

        end_time = datetime.now()
        generation_time = (end_time - start_time).total_seconds()
        self.logger.info(f"排班方案生成完成，耗时 {generation_time:.2f} 秒")

        return schedule

    def preprocess_historical_data(
        self, historical_schedules, caregivers, elders, shift_config
    ):
        """预处理历史排班数据

        Args:
            historical_schedules: 历史排班数据
            caregivers: 护理员列表
            elders: 老人列表
            shift_config: 班次配置

        Returns:
            状态-动作-奖励-下一状态序列
        """
        history = []

        for i in range(len(historical_schedules) - 1):
            current_schedule = historical_schedules[i]
            next_schedule = historical_schedules[i + 1]

            # 提取状态
            state = self.schedule_to_state(
                current_schedule, caregivers, elders, shift_config
            )
            next_state = self.schedule_to_state(
                next_schedule, caregivers, elders, shift_config
            )

            # 提取动作
            action = self.extract_action(current_schedule, next_schedule)

            # 计算奖励
            reward = self.calculate_reward(
                next_schedule, caregivers, elders, shift_config
            )

            history.append((state, action, reward, next_state))

        return history

    def schedule_to_state(self, schedule, caregivers, elders, shift_config):
        """将排班方案转换为状态

        Args:
            schedule: 排班方案
            caregivers: 护理员列表
            elders: 老人列表
            shift_config: 班次配置

        Returns:
            状态表示
        """
        # 状态表示为一个元组，包含以下信息：
        # 1. 每个护理员的班次分配
        # 2. 每个老人的照护计划

        # 提取护理员班次分配
        caregiver_shifts = {}
        if "shift_assignments" in schedule:
            for assignment in schedule["shift_assignments"]:
                caregiver_id = assignment.caregiver_id
                date_key = str(assignment.shift_date)
                shift_type = assignment.shift_type

                if caregiver_id not in caregiver_shifts:
                    caregiver_shifts[caregiver_id] = {}

                if date_key not in caregiver_shifts[caregiver_id]:
                    caregiver_shifts[caregiver_id][date_key] = []

                caregiver_shifts[caregiver_id][date_key].append(shift_type)

        # 提取老人照护计划
        elder_care_plans = {}
        if "schedule_results" in schedule:
            for result in schedule["schedule_results"]:
                elder_id = result.elder_id
                date_key = str(result.date)

                if elder_id not in elder_care_plans:
                    elder_care_plans[elder_id] = {}

                elder_care_plans[elder_id][date_key] = result.slots_count

        # 将字典转换为不可变的元组
        caregiver_shifts_tuple = tuple(
            sorted(
                (
                    caregiver_id,
                    tuple(
                        sorted(
                            (date, tuple(sorted(shifts)))
                            for date, shifts in dates.items()
                        )
                    ),
                )
                for caregiver_id, dates in caregiver_shifts.items()
            )
        )

        elder_care_plans_tuple = tuple(
            sorted(
                (
                    elder_id,
                    tuple(sorted((date, count) for date, count in dates.items())),
                )
                for elder_id, dates in elder_care_plans.items()
            )
        )

        return (caregiver_shifts_tuple, elder_care_plans_tuple)

    def state_to_schedule(self, state, caregivers, elders, shift_config):
        """将状态转换为排班方案

        Args:
            state: 状态
            caregivers: 护理员列表
            elders: 老人列表
            shift_config: 班次配置

        Returns:
            排班方案
        """
        # 这里需要根据实际的排班方案数据结构进行实现
        # 简单起见，假设排班方案是一个字典
        schedule = {
            "shift_assignments": [],
            "schedule_results": [],
            "time_slot_assignments": [],
        }

        caregiver_shifts_tuple, elder_care_plans_tuple = state

        # 转换护理员班次分配
        for caregiver_id, dates_tuple in caregiver_shifts_tuple:
            for date_str, shifts_tuple in dates_tuple:
                for shift_type in shifts_tuple:
                    # 创建班次分配对象
                    assignment = type(
                        "ShiftAssignment",
                        (),
                        {
                            "caregiver_id": caregiver_id,
                            "shift_date": datetime.strptime(
                                date_str, "%Y-%m-%d"
                            ).date(),
                            "shift_type": shift_type,
                        },
                    )
                    schedule["shift_assignments"].append(assignment)

        # 转换老人照护计划
        for elder_id, dates_tuple in elder_care_plans_tuple:
            for date_str, slots_count in dates_tuple:
                # 创建照护计划对象
                result = type(
                    "ScheduleResult",
                    (),
                    {
                        "elder_id": elder_id,
                        "date": datetime.strptime(date_str, "%Y-%m-%d").date(),
                        "slots_count": slots_count,
                    },
                )
                schedule["schedule_results"].append(result)

        return schedule

    def extract_action(self, current_schedule, next_schedule):
        """提取从当前排班方案到下一个排班方案的动作

        Args:
            current_schedule: 当前排班方案
            next_schedule: 下一个排班方案

        Returns:
            动作表示
        """
        # 动作表示为一个元组，包含以下信息：
        # 1. 添加的班次分配
        # 2. 删除的班次分配
        # 3. 修改的班次分配

        # 提取当前排班方案的班次分配
        current_assignments = {}
        if "shift_assignments" in current_schedule:
            for assignment in current_schedule["shift_assignments"]:
                key = (
                    assignment.caregiver_id,
                    str(assignment.shift_date),
                    assignment.shift_type,
                )
                current_assignments[key] = assignment

        # 提取下一个排班方案的班次分配
        next_assignments = {}
        if "shift_assignments" in next_schedule:
            for assignment in next_schedule["shift_assignments"]:
                key = (
                    assignment.caregiver_id,
                    str(assignment.shift_date),
                    assignment.shift_type,
                )
                next_assignments[key] = assignment

        # 计算添加、删除和修改的班次分配
        added = tuple(
            sorted(set(next_assignments.keys()) - set(current_assignments.keys()))
        )
        deleted = tuple(
            sorted(set(current_assignments.keys()) - set(next_assignments.keys()))
        )

        return (added, deleted)

    def calculate_reward(self, schedule, caregivers, elders, shift_config):
        """计算排班方案的奖励

        Args:
            schedule: 排班方案
            caregivers: 护理员列表
            elders: 老人列表
            shift_config: 班次配置

        Returns:
            奖励值
        """
        reward = 0

        # 1. 班次覆盖奖励
        if "shift_assignments" in schedule and schedule["shift_assignments"]:
            # 按日期和班次类型分组
            shifts_by_date_type = {}
            for assignment in schedule["shift_assignments"]:
                date_key = str(assignment.shift_date)
                if date_key not in shifts_by_date_type:
                    shifts_by_date_type[date_key] = {}

                shift_type = assignment.shift_type
                if shift_type not in shifts_by_date_type[date_key]:
                    shifts_by_date_type[date_key][shift_type] = 0

                shifts_by_date_type[date_key][shift_type] += 1

            # 计算班次覆盖率
            total_coverage = 0
            total_days = len(shifts_by_date_type)
            total_shift_types = len(shift_config)

            if total_days > 0 and total_shift_types > 0:
                for date_key, shifts in shifts_by_date_type.items():
                    coverage = len(shifts) / total_shift_types
                    total_coverage += coverage

                average_coverage = total_coverage / total_days
                reward += average_coverage * 100  # 满分100

        # 2. 工作负载平衡奖励
        if "shift_assignments" in schedule and schedule["shift_assignments"]:
            # 统计每个护理员的班次数
            caregiver_shifts = {}
            for assignment in schedule["shift_assignments"]:
                caregiver_id = assignment.caregiver_id
                if caregiver_id not in caregiver_shifts:
                    caregiver_shifts[caregiver_id] = 0

                caregiver_shifts[caregiver_id] += 1

            # 计算标准差
            if caregiver_shifts:
                shift_counts = list(caregiver_shifts.values())
                mean_shifts = sum(shift_counts) / len(shift_counts)

                if mean_shifts > 0:
                    std_dev = (
                        sum((count - mean_shifts) ** 2 for count in shift_counts)
                        / len(shift_counts)
                    ) ** 0.5

                    # 标准差越小，平衡度越高，奖励越大
                    balance_score = 50 / (1 + std_dev)  # 最高分值50
                    reward += balance_score

        # 3. 照护计划完整性奖励
        if "schedule_results" in schedule and schedule["schedule_results"]:
            total_slots = 0
            expected_slots = 0

            for result in schedule["schedule_results"]:
                total_slots += result.slots_count
                expected_slots += 48  # 每个老人每天应该有48个时间段

            if expected_slots > 0:
                completeness = total_slots / expected_slots
                reward += completeness * 50  # 最高分值50

        return reward

    def select_action(self, state):
        """选择动作

        Args:
            state: 当前状态

        Returns:
            选择的动作
        """
        # 探索：随机选择动作
        if random.random() < self.exploration_rate:
            # 生成一个随机动作
            return self.generate_random_action(state)

        # 利用：选择Q值最大的动作
        if state in self.q_table:
            actions = list(self.q_table[state].keys())
            if actions:
                return max(actions, key=lambda a: self.q_table[state][a])

        # 如果状态不在Q值表中，或者没有可用的动作，则生成一个随机动作
        return self.generate_random_action(state)

    def generate_random_action(self, state):
        """生成随机动作

        Args:
            state: 当前状态

        Returns:
            随机动作
        """
        # 简单起见，生成一个空动作
        return ((), ())

    def take_action(self, state, action, caregivers, elders, shift_config):
        """执行动作，获取下一个状态和奖励

        Args:
            state: 当前状态
            action: 选择的动作
            caregivers: 护理员列表
            elders: 老人列表
            shift_config: 班次配置

        Returns:
            下一个状态和奖励
        """
        # 将状态转换为排班方案
        schedule = self.state_to_schedule(state, caregivers, elders, shift_config)

        # 执行动作，修改排班方案
        added, deleted = action

        # 删除班次分配
        if "shift_assignments" in schedule:
            schedule["shift_assignments"] = [
                assignment
                for assignment in schedule["shift_assignments"]
                if (
                    assignment.caregiver_id,
                    str(assignment.shift_date),
                    assignment.shift_type,
                )
                not in deleted
            ]

        # 添加班次分配
        for caregiver_id, date_str, shift_type in added:
            # 创建班次分配对象
            assignment = type(
                "ShiftAssignment",
                (),
                {
                    "caregiver_id": caregiver_id,
                    "shift_date": datetime.strptime(date_str, "%Y-%m-%d").date(),
                    "shift_type": shift_type,
                },
            )
            schedule["shift_assignments"].append(assignment)

        # 更新照护计划
        self.update_care_plans(schedule, caregivers, elders, shift_config)

        # 计算奖励
        reward = self.calculate_reward(schedule, caregivers, elders, shift_config)

        # 将排班方案转换为下一个状态
        next_state = self.schedule_to_state(schedule, caregivers, elders, shift_config)

        return next_state, reward

    def update_care_plans(self, schedule, caregivers, elders, shift_config):
        """更新照护计划

        Args:
            schedule: 排班方案
            caregivers: 护理员列表
            elders: 老人列表
            shift_config: 班次配置
        """
        # 这里需要根据实际的排班方案数据结构进行实现
        # 简单起见，假设每个老人每天有48个时间段
        if "schedule_results" not in schedule:
            schedule["schedule_results"] = []

        # 获取所有日期
        dates = set()
        if "shift_assignments" in schedule:
            for assignment in schedule["shift_assignments"]:
                dates.add(assignment.shift_date)

        # 为每个老人每天创建照护计划
        for elder in elders:
            for date in dates:
                # 检查是否已经存在照护计划
                existing = False
                for result in schedule["schedule_results"]:
                    if result.elder_id == elder.id and result.date == date:
                        existing = True
                        break

                if not existing:
                    # 创建照护计划对象
                    result = type(
                        "ScheduleResult",
                        (),
                        {
                            "elder_id": elder.id,
                            "date": date,
                            "slots_count": 48,  # 每天48个时间段
                        },
                    )
                    schedule["schedule_results"].append(result)

    def update_q_value(self, state, action, reward, next_state):
        """更新Q值

        Args:
            state: 当前状态
            action: 选择的动作
            reward: 获得的奖励
            next_state: 下一个状态
        """
        # 获取当前Q值
        current_q = self.q_table[state][action]

        # 获取下一个状态的最大Q值
        max_next_q = 0
        if next_state in self.q_table:
            actions = list(self.q_table[next_state].keys())
            if actions:
                max_next_q = max(self.q_table[next_state][a] for a in actions)

        # 更新Q值
        new_q = current_q + self.learning_rate * (
            reward + self.discount_factor * max_next_q - current_q
        )
        self.q_table[state][action] = new_q

    def is_terminal_state(self, state, caregivers, elders, shift_config):
        """判断是否为终止状态

        Args:
            state: 当前状态
            caregivers: 护理员列表
            elders: 老人列表
            shift_config: 班次配置

        Returns:
            是否为终止状态
        """
        # 将状态转换为排班方案
        schedule = self.state_to_schedule(state, caregivers, elders, shift_config)

        # 计算奖励
        reward = self.calculate_reward(schedule, caregivers, elders, shift_config)

        # 如果奖励足够高，则认为是终止状态
        return reward >= 180  # 满分200
